import logging

from fastapi import APIRouter, Body, Query
from tortoise.expressions import Q

from app.controllers.model import model_controller
from app.schemas.base import Fail, Success, SuccessExtra
from app.schemas.models import *

logger = logging.getLogger(__name__)

router = APIRouter()


@router.get("/list", summary="查看模型列表")
async def list_model(
    page: int = Query(1, description="页码"),
    page_size: int = Query(10, description="每页数量"),
    name: str = Query("", description="名称，用于搜索"),
    frame: str = Query("", description="模型框架，如YOLOv8"),
    task: str = Query("", description="任务类型 是属于分类还是其他的"),
    is_quad: bool = Query(True, description="任务类型 是属于分类还是其他的"),
):
    q = Q()
    if name:
        q &= Q(name__contains=name)
    if frame:
        q &= Q(frame__contains=frame)
    if task:
        q &= Q(task__contains=task)
    if is_quad:
        q &= Q(is_quad=is_quad)
    total, model_objs = await model_controller.list(page=page, page_size=page_size, search=q)
    data = [await obj.to_dict(m2m=True, exclude_fields=[]) for obj in model_objs]
    return SuccessExtra(data=data, total=total, page=page, page_size=page_size)


@router.get("/get", summary="查看模型")
async def get_model(
    model_id: int = Query(..., description="模型ID"),
):
    model_obj = await model_controller.get(id=model_id)
    model_dict = await model_obj.to_dict(exclude_fields=["password"])
    return Success(data=model_dict)


@router.post("/create", summary="创建模型")
async def create_model(
    model_in: ModelCreate,
):
    model = await model_controller.get_by_username(model_in.name)
    if model:
        return Fail(code=400, msg="系统中已存在带有此名称的模型。")
    new_model = await model_controller.create_model(obj_in=model_in)
    model_dict = await new_model.to_dict(exclude_fields=[])
    return Success(msg="Created Successfully",data=model_dict)


@router.post("/update", summary="更新模型")
async def update_model(
    model_in: ModelUpdate,
):
    model = await model_controller.update(id=model_in.id, obj_in=model_in)
    model_dict = await model.to_dict(exclude_fields=[])
    return Success(msg="Updated Successfully",data=model_dict)


@router.delete("/delete", summary="删除模型")
async def delete_model(
    model_id: int = Query(..., description="模型ID"),
):
    await model_controller.remove(id=model_id)
    return Success(msg="Deleted Successfully")

